#!/usr/bin/env python

import lattice
import math

def exportTreeToPDE(filename,labels,dim=200,rmax=0):
    """writes a pde file for processing.org"""
    if rmax == 0:
        rmax = math.floor(math.log(max(labels))/math.log(2))
    offset = dim/2
    rscaling = dim/2 / rmax

    f = open(filename, 'w')
    print >> f, "size(",dim,",",dim,"); smooth(); background(127); fill(0);"
    print >> f, "ellipse(",dim/2,",",dim/2,",",dim,",",dim,");"
    print >> f, "stroke(255);"

    for k in labels:
        if 2*k in labels:
            k1c = lattice.convertToXYGrid(lattice.calcPolarCoords(k),rscaling,offset)
            k2c = lattice.convertToXYGrid(lattice.calcPolarCoords(2*k),rscaling,offset)
            print >> f, "line(",k1c[0],",",k1c[1],",",k2c[0],",",k2c[1],");"
        if 2*k+1 in labels:
            k1c = lattice.convertToXYGrid(lattice.calcPolarCoords(k),rscaling,offset)
            k2c = lattice.convertToXYGrid(lattice.calcPolarCoords(2*k+1),rscaling,offset)
            print >> f, "line(",k1c[0],",",k1c[1],",",k2c[0],",",k2c[1],");"
    f.close()

def exportPopulationToPDE(filename, popSize, popLabels, rmax, dim=200):
    """writes a pde file for processing.org"""
    offset = dim/2
    rscaling = dim/2 / rmax

    f = open(filename, 'w')
    print >> f, "size(",dim,",",dim,"); smooth(); background(127); fill(0);"
    print >> f, "ellipse(",dim/2,",",dim/2,",",dim,",",dim,");"
    print >> f, "stroke(255);"

    for k in popLabels:
        if 2*k in popLabels:
            k1c = lattice.convertToXYGrid(lattice.calcPolarCoords(k),rscaling,offset)
            k2c = lattice.convertToXYGrid(lattice.calcPolarCoords(2*k),rscaling,offset)
            print >> f, "stroke(",245*popLabels[k]/popSize+10,");"
            print >> f, "line(",k1c[0],",",k1c[1],",",k2c[0],",",k2c[1],");"
        if 2*k+1 in popLabels:
            k1c = lattice.convertToXYGrid(lattice.calcPolarCoords(k),rscaling,offset)
            k2c = lattice.convertToXYGrid(lattice.calcPolarCoords(2*k+1),rscaling,offset)
            print >> f, "stroke(",245*popLabels[k]/popSize+10,");"
            print >> f, "line(",k1c[0],",",k1c[1],",",k2c[0],",",k2c[1],");"
    f.close()

    pass


if __name__ == '__main__':
    fs = ['+','-','*','%']
    ts = [1,2]

    depth = 10

    pop = []
    indivs = 10
    import trees

    for i in range(indivs):
        gp = trees.generateRandomExpr(fs, ts, depth, 'grow')
        pop = pop +[gp]

    oc = lattice.parsePopulationToLattice(pop)
    exportPopulationToPDE('alpha.pde', indivs, oc, depth, 400)


